msg_tool\scripts\artemis/
asb.rs

1//! Artemis Engine ASB file (.asb)
2use crate::ext::io::*;
3use crate::scripts::base::*;
4use crate::types::*;
5use crate::utils::encoding::*;
6use crate::utils::escape::*;
7use anyhow::Result;
8use serde::{Deserialize, Serialize};
9use std::collections::BTreeMap;
10use std::io::{Read, Write};
11use std::ops::Index;
12use unicode_segmentation::UnicodeSegmentation;
13
14#[derive(Debug)]
15/// The builder for Artemis ASB scripts.
16pub struct ArtemisAsbBuilder {}
17
18impl ArtemisAsbBuilder {
19    /// Creates a new instance of `ArtemisAsbBuilder`.
20    pub fn new() -> Self {
21        ArtemisAsbBuilder {}
22    }
23}
24
25impl ScriptBuilder for ArtemisAsbBuilder {
26    fn default_encoding(&self) -> Encoding {
27        Encoding::Utf8
28    }
29
30    fn build_script(
31        &self,
32        buf: Vec<u8>,
33        _filename: &str,
34        encoding: Encoding,
35        _archive_encoding: Encoding,
36        config: &ExtraConfig,
37        _archive: Option<&Box<dyn Script>>,
38    ) -> Result<Box<dyn Script>> {
39        Ok(Box::new(Asb::new(buf, encoding, config)?))
40    }
41
42    fn extensions(&self) -> &'static [&'static str] {
43        &["asb"]
44    }
45
46    fn script_type(&self) -> &'static ScriptType {
47        &ScriptType::ArtemisAsb
48    }
49
50    fn is_this_format(&self, _filename: &str, buf: &[u8], buf_len: usize) -> Option<u8> {
51        if buf_len >= 5 && buf.starts_with(b"ASB\0\0") {
52            return Some(20);
53        }
54        None
55    }
56
57    fn can_create_file(&self) -> bool {
58        true
59    }
60
61    fn create_file<'a>(
62        &'a self,
63        filename: &'a str,
64        writer: Box<dyn WriteSeek + 'a>,
65        encoding: Encoding,
66        file_encoding: Encoding,
67        config: &ExtraConfig,
68    ) -> Result<()> {
69        create_file(
70            filename,
71            writer,
72            encoding,
73            file_encoding,
74            config.custom_yaml,
75        )
76    }
77}
78
79fn escape_text(s: &str) -> String {
80    let mut escaped = String::with_capacity(s.len());
81    for c in s.chars() {
82        match c {
83            '&' => escaped.push_str("&amp;"),
84            '<' => escaped.push_str("&lt;"),
85            _ => escaped.push(c),
86        }
87    }
88    escaped
89}
90
91#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
92struct Command {
93    pub name: String,
94    pub line_number: u32,
95    pub attributes: BTreeMap<String, String>,
96}
97
98impl Command {
99    pub fn new(name: String, line_number: u32) -> Self {
100        Command {
101            name,
102            line_number,
103            attributes: BTreeMap::new(),
104        }
105    }
106
107    pub fn to_xml(&self) -> String {
108        let mut xml = format!("<{}", self.name);
109        for (key, value) in &self.attributes {
110            xml.push_str(&format!(" {}=\"{}\"", key, escape_xml_text_value(value)));
111        }
112        xml.push('>');
113        xml
114    }
115}
116
117impl<'a> Index<&'a str> for Command {
118    type Output = str;
119    fn index(&self, key: &'a str) -> &Self::Output {
120        self.attributes.get(key).map_or("", |s| s.as_str())
121    }
122}
123
124impl<'a> Index<&'a String> for Command {
125    type Output = str;
126    fn index(&self, key: &'a String) -> &Self::Output {
127        self.attributes.get(key).map_or("", |s| s.as_str())
128    }
129}
130
131impl Index<String> for Command {
132    type Output = str;
133    fn index(&self, key: String) -> &Self::Output {
134        self.attributes.get(&key).map_or("", |s| s.as_str())
135    }
136}
137
138#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
139#[serde(untagged)]
140enum Item {
141    Command(Command),
142    Label(String),
143}
144
145impl Item {
146    pub fn is_command(&self) -> bool {
147        matches!(self, Item::Command(_))
148    }
149
150    pub fn is_command_name(&self, name: &str) -> bool {
151        if let Item::Command(cmd) = self {
152            cmd.name == name
153        } else {
154            false
155        }
156    }
157}
158
159trait CustomReadFn {
160    fn read_string(&mut self, encoding: Encoding) -> Result<String>;
161    fn read_item(&mut self, encoding: Encoding) -> Result<Item>;
162}
163
164impl<T: Read> CustomReadFn for T {
165    fn read_string(&mut self, encoding: Encoding) -> Result<String> {
166        let len = self.read_u32()?;
167        let data = self.read_exact_vec(len as usize)?;
168        if self.read_u8()? != 0 {
169            return Err(anyhow::anyhow!("String not null-terminated"));
170        }
171        let s = decode_to_string(encoding, &data, true)?;
172        Ok(s)
173    }
174
175    fn read_item(&mut self, encoding: Encoding) -> Result<Item> {
176        let typ = self.read_u32()?;
177        match typ {
178            0 => {
179                let name = self.read_string(encoding)?;
180                let line_number = self.read_u32()?;
181                let mut command = Command::new(name, line_number);
182                let attr_count = self.read_u32()?;
183                for _ in 0..attr_count {
184                    let key = self.read_string(encoding)?;
185                    let value = self.read_string(encoding)?;
186                    command.attributes.insert(key, value);
187                }
188                Ok(Item::Command(command))
189            }
190            1 => {
191                let label = self.read_string(encoding)?;
192                Ok(Item::Label(label))
193            }
194            _ => {
195                return Err(anyhow::anyhow!("Unknown item type: {}", typ));
196            }
197        }
198    }
199}
200
201trait CustomWriteFn {
202    fn write_string(&mut self, s: &str, encoding: Encoding) -> Result<()>;
203    fn write_item(&mut self, item: &Item, encoding: Encoding) -> Result<()>;
204}
205
206impl<T: Write> CustomWriteFn for T {
207    fn write_string(&mut self, s: &str, encoding: Encoding) -> Result<()> {
208        let data = encode_string(encoding, s, false)?;
209        self.write_u32(data.len() as u32)?;
210        self.write_all(&data)?;
211        self.write_u8(0)?; // Null-terminated
212        Ok(())
213    }
214
215    fn write_item(&mut self, item: &Item, encoding: Encoding) -> Result<()> {
216        match item {
217            Item::Command(cmd) => {
218                self.write_u32(0)?; // Type 0 for Command
219                self.write_string(&cmd.name, encoding)?;
220                self.write_u32(cmd.line_number)?;
221                self.write_u32(cmd.attributes.len() as u32)?;
222                for (key, value) in &cmd.attributes {
223                    self.write_string(key, encoding)?;
224                    self.write_string(value, encoding)?;
225                }
226            }
227            Item::Label(label) => {
228                self.write_u32(1)?; // Type 1 for Label
229                self.write_string(label, encoding)?;
230            }
231        }
232        Ok(())
233    }
234}
235
236struct TextParser<'a> {
237    items: Vec<Item>,
238    text: Vec<&'a str>,
239    pos: usize,
240    len: usize,
241    hcls_index: usize,
242}
243
244impl<'a> TextParser<'a> {
245    pub fn new(str: &'a str, hcls_index: usize) -> Self {
246        let text: Vec<&'a str> = UnicodeSegmentation::graphemes(str, true).collect();
247        let len = text.len();
248        TextParser {
249            items: Vec::new(),
250            text,
251            pos: 0,
252            len,
253            hcls_index,
254        }
255    }
256
257    pub fn parse(mut self) -> Result<Vec<Item>> {
258        while let Some(c) = self.peek() {
259            match c {
260                "<" => {
261                    self.parse_tag()?;
262                }
263                _ => {
264                    let mut text = String::new();
265                    self.eat_char();
266                    text.push_str(c);
267                    while let Some(b) = self.peek() {
268                        if b == "<" {
269                            break;
270                        }
271                        text.push_str(b);
272                        self.eat_char();
273                    }
274                    if !text.is_empty() {
275                        self.items.push(Item::Command(Command {
276                            name: "print".to_string(),
277                            line_number: 0,
278                            attributes: [("data".to_string(), unescape_xml(&text))].into(),
279                        }))
280                    }
281                }
282            }
283        }
284        let mut hcls = Command::new("hcls".to_string(), 0);
285        hcls.attributes
286            .insert("0".to_string(), self.hcls_index.to_string());
287        self.items.push(Item::Command(hcls));
288        Ok(self.items)
289    }
290
291    fn parse_tag(&mut self) -> Result<()> {
292        self.parse_indent("<")?;
293        let key = self.parse_key()?;
294        self.erase_whitespace();
295        let mut cmd = Command::new(key, 0);
296        loop {
297            let c = self.peek().ok_or(self.error2("Unexpected eof"))?;
298            match c {
299                ">" => {
300                    self.eat_char();
301                    break;
302                }
303                " " => {
304                    self.eat_char();
305                    continue;
306                }
307                _ => {
308                    let key = self.parse_key()?;
309                    self.parse_indent("=")?;
310                    let value = self.parse_str()?;
311                    cmd.attributes.insert(key, value);
312                }
313            }
314        }
315        self.items.push(Item::Command(cmd));
316        Ok(())
317    }
318
319    fn parse_key(&mut self) -> Result<String> {
320        self.erase_whitespace();
321        let mut key = String::new();
322        while let Some(c) = self.peek() {
323            if c == "=" || c == " " || c == ">" {
324                break;
325            }
326            key.push_str(c);
327            self.eat_char();
328        }
329        if key.is_empty() {
330            return self.error("Expected key, but found nothing");
331        }
332        Ok(key)
333    }
334
335    fn parse_str(&mut self) -> Result<String> {
336        self.erase_whitespace();
337        self.parse_indent("\"")?;
338        let mut text = String::new();
339        loop {
340            match self.next().ok_or(self.error2("Unexpected eof"))? {
341                "\"" => {
342                    break;
343                }
344                t => {
345                    text.push_str(t);
346                }
347            }
348        }
349        Ok(unescape_xml(&text))
350    }
351
352    fn erase_whitespace(&mut self) {
353        while let Some(c) = self.peek() {
354            if c == " " {
355                self.eat_char();
356            } else {
357                break;
358            }
359        }
360    }
361
362    fn parse_indent(&mut self, indent: &str) -> Result<()> {
363        for ident in indent.graphemes(true) {
364            match self.next() {
365                Some(c) => {
366                    if c != ident {
367                        return self.error("Unexpected indent");
368                    }
369                }
370                None => return self.error("Unexpected eof"),
371            }
372        }
373        Ok(())
374    }
375
376    fn eat_char(&mut self) {
377        if self.pos < self.len {
378            self.pos += 1;
379        }
380    }
381
382    fn next(&mut self) -> Option<&'a str> {
383        if self.pos < self.len {
384            let item = self.text[self.pos];
385            self.pos += 1;
386            Some(item)
387        } else {
388            None
389        }
390    }
391
392    fn peek(&self) -> Option<&'a str> {
393        if self.pos < self.len {
394            Some(self.text[self.pos])
395        } else {
396            None
397        }
398    }
399
400    fn error2<T>(&self, msg: T) -> anyhow::Error
401    where
402        T: std::fmt::Display,
403    {
404        anyhow::anyhow!("Failed to parse at position {}: {}", self.pos, msg)
405    }
406
407    fn error<T, A>(&self, msg: T) -> Result<A>
408    where
409        T: std::fmt::Display,
410    {
411        Err(anyhow::anyhow!(
412            "Failed to parse at position {}: {}",
413            self.pos,
414            msg
415        ))
416    }
417}
418
419#[derive(Debug)]
420/// The Artemis ASB script.
421pub struct Asb {
422    items: Vec<Item>,
423    custom_yaml: bool,
424}
425
426impl Asb {
427    /// Creates a new Artemis ASB script from the given buffer.
428    ///
429    /// * `buf` - The buffer containing the ASB data.
430    /// * `encoding` - The encoding used for the ASB data.
431    /// * `config` - Extra configuration options.
432    pub fn new(buf: Vec<u8>, encoding: Encoding, config: &ExtraConfig) -> Result<Self> {
433        let mut data = MemReader::new(buf);
434        let mut magic = [0; 5];
435        data.read_exact(&mut magic)?;
436        if &magic != b"ASB\0\0" {
437            return Err(anyhow::anyhow!("Invalid ASB magic number: {:?}", magic));
438        }
439        let nums = data.read_u32()?;
440        let mut items = Vec::with_capacity(nums as usize);
441        for _ in 0..nums {
442            items.push(data.read_item(encoding)?);
443        }
444        Ok(Asb {
445            items,
446            custom_yaml: config.custom_yaml,
447        })
448    }
449}
450
451impl Script for Asb {
452    fn default_output_script_type(&self) -> OutputScriptType {
453        OutputScriptType::Json
454    }
455
456    fn is_output_supported(&self, _: OutputScriptType) -> bool {
457        true
458    }
459
460    fn default_format_type(&self) -> FormatOptions {
461        FormatOptions::None
462    }
463
464    fn custom_output_extension<'a>(&'a self) -> &'a str {
465        if self.custom_yaml { "yaml" } else { "json" }
466    }
467
468    fn extract_messages(&self) -> Result<Vec<Message>> {
469        let mut messages = Vec::new();
470        let mut name = None;
471        let mut cur_mes = String::new();
472        let mut in_print = false;
473        for item in self.items.iter() {
474            if in_print {
475                if let Item::Command(cmd) = item {
476                    match cmd.name.as_str() {
477                        "hcls" => {
478                            in_print = false;
479                            messages.push(Message {
480                                name: name.take(),
481                                message: cur_mes,
482                            });
483                            cur_mes = String::new();
484                        }
485                        "print" => {
486                            cur_mes.push_str(&escape_text(&cmd["data"]));
487                        }
488                        "rt" => {
489                            cur_mes.push('\n');
490                        }
491                        _ => {
492                            cur_mes.push_str(&cmd.to_xml());
493                        }
494                    }
495                    continue;
496                }
497            }
498            if let Item::Command(cmd) = item {
499                match cmd.name.as_str() {
500                    "print" => {
501                        cur_mes.push_str(&escape_text(&cmd["data"]));
502                        in_print = true;
503                    }
504                    "name" => {
505                        let v = (cmd.attributes.len() - 1).to_string();
506                        name = Some(cmd[v].to_owned());
507                    }
508                    "sel_text" => {
509                        let t = &cmd["text"];
510                        if !t.is_empty() {
511                            messages.push(Message {
512                                name: None,
513                                message: t.to_owned(),
514                            });
515                        }
516                    }
517                    "RegisterTextToHistory" => {
518                        let t = &cmd["1"];
519                        if !t.is_empty() {
520                            messages.push(Message {
521                                name: None,
522                                message: t.to_owned(),
523                            });
524                        }
525                    }
526                    _ => {}
527                }
528            }
529        }
530        if !cur_mes.is_empty() {
531            messages.push(Message {
532                name: name.take(),
533                message: cur_mes,
534            });
535        }
536        Ok(messages)
537    }
538
539    fn import_messages<'a>(
540        &'a self,
541        messages: Vec<Message>,
542        mut file: Box<dyn WriteSeek + 'a>,
543        _filename: &str,
544        encoding: Encoding,
545        replacement: Option<&'a ReplacementTable>,
546    ) -> Result<()> {
547        file.write_all(b"ASB\0\0")?;
548        let mut items = self.items.clone();
549        let mut name_index = None;
550        let mut mes_index = 0;
551        let mut item_index = 0;
552        let mut print_index = None;
553        let mut hcls_index = 1;
554        while item_index < items.len() {
555            if let Some(print_ind) = print_index.clone() {
556                if items[item_index].is_command_name("hcls") {
557                    let message = messages
558                        .get(mes_index)
559                        .ok_or(anyhow::anyhow!("Not enough messages."))?;
560                    if let Some(name_index) = name_index.take() {
561                        let mut name = match &message.name {
562                            Some(name) => name.to_owned(),
563                            None => return Err(anyhow::anyhow!("Message without name.")),
564                        };
565                        if let Some(replacement) = replacement {
566                            for (k, v) in &replacement.map {
567                                name = name.replace(k, v);
568                            }
569                        }
570                        if let Item::Command(cmd) = &mut items[name_index] {
571                            if cmd.attributes.len() > 1 {
572                                cmd.attributes
573                                    .insert(format!("{}", cmd.attributes.len() - 1), name);
574                            } else {
575                                let oname = cmd
576                                    .attributes
577                                    .get("0")
578                                    .ok_or(anyhow::anyhow!("No name attribute found."))?;
579                                if oname != &name {
580                                    cmd.attributes.insert("1".to_string(), name);
581                                }
582                            }
583                        }
584                    }
585                    let mut m = message.message.clone();
586                    if let Some(replacement) = replacement {
587                        for (k, v) in &replacement.map {
588                            m = m.replace(k, v);
589                        }
590                    }
591                    let new_cmds = TextParser::new(&m.replace("\n", "<rt>"), hcls_index).parse()?;
592                    hcls_index += 1;
593                    let new_cmds_len = new_cmds.len();
594                    items.splice(print_ind..=item_index, new_cmds);
595                    print_index = None;
596                    item_index = print_ind + new_cmds_len;
597                    mes_index += 1;
598                    continue;
599                } else if items[item_index].is_command() {
600                    item_index += 1;
601                    continue;
602                }
603            }
604            if let Item::Command(cmd) = &mut items[item_index] {
605                match cmd.name.as_str() {
606                    "print" => {
607                        print_index = Some(item_index);
608                    }
609                    "name" => {
610                        name_index = Some(item_index);
611                    }
612                    "sel_text" => {
613                        let message = messages
614                            .get(mes_index)
615                            .ok_or(anyhow::anyhow!("Not enough messages."))?;
616                        let mut m = message.message.clone();
617                        if let Some(replacement) = replacement {
618                            for (k, v) in &replacement.map {
619                                m = m.replace(k, v);
620                            }
621                        }
622                        cmd.attributes.insert("text".to_string(), m);
623                        mes_index += 1;
624                    }
625                    "RegisterTextToHistory" => {
626                        let message = messages
627                            .get(mes_index)
628                            .ok_or(anyhow::anyhow!("Not enough messages."))?;
629                        let mut m = message.message.clone();
630                        if let Some(replacement) = replacement {
631                            for (k, v) in &replacement.map {
632                                m = m.replace(k, v);
633                            }
634                        }
635                        cmd.attributes.insert("1".to_string(), m);
636                        mes_index += 1;
637                    }
638                    _ => {}
639                }
640            }
641            item_index += 1;
642        }
643        if mes_index != messages.len() {
644            return Err(anyhow::anyhow!(
645                "Not all messages were processed, expected {}, got {}",
646                messages.len(),
647                mes_index
648            ));
649        }
650        file.write_u32(items.len() as u32)?;
651        for item in items {
652            file.write_item(&item, encoding)?;
653        }
654        file.flush()?;
655        Ok(())
656    }
657
658    fn custom_export(&self, filename: &std::path::Path, encoding: Encoding) -> Result<()> {
659        let s = if self.custom_yaml {
660            serde_yaml_ng::to_string(&self.items)?
661        } else {
662            serde_json::to_string_pretty(&self.items)?
663        };
664        let s = encode_string(encoding, &s, false)?;
665        let mut file = std::fs::File::create(filename)?;
666        file.write_all(&s)?;
667        Ok(())
668    }
669
670    fn custom_import<'a>(
671        &'a self,
672        custom_filename: &'a str,
673        file: Box<dyn WriteSeek + 'a>,
674        encoding: Encoding,
675        output_encoding: Encoding,
676    ) -> Result<()> {
677        create_file(
678            custom_filename,
679            file,
680            encoding,
681            output_encoding,
682            self.custom_yaml,
683        )
684    }
685}
686
687/// Creates a new ASB file.
688///
689/// * `custom_filename` - The path ot the input file.
690/// * `writer` - The writer to write the ASB script.
691/// * `encoding` - The encoding used for the ASB script.
692/// * `output_encoding` - The encoding used for the input file.
693/// * `yaml` - Whether to use YAML format instead of JSON for the input file.
694pub fn create_file<'a>(
695    custom_filename: &'a str,
696    mut writer: Box<dyn WriteSeek + 'a>,
697    encoding: Encoding,
698    output_encoding: Encoding,
699    yaml: bool,
700) -> Result<()> {
701    let f = crate::utils::files::read_file(custom_filename)?;
702    let s = decode_to_string(output_encoding, &f, true)?;
703    let items: Vec<Item> = if yaml {
704        serde_yaml_ng::from_str(&s)?
705    } else {
706        serde_json::from_str(&s)?
707    };
708    writer.write_all(b"ASB\0\0")?;
709    writer.write_u32(items.len() as u32)?;
710    for item in items {
711        writer.write_item(&item, encoding)?;
712    }
713    Ok(())
714}
715
716#[test]
717fn test_parse() {
718    let text = "Hello &lt; &amp; World!<tag><tags x=\"123\"><name 0=\"Ok\">Test";
719    let parser = TextParser::new(text, 1);
720    let items = parser.parse().unwrap();
721    assert_eq!(
722        items,
723        vec![
724            Item::Command(Command {
725                name: "print".to_string(),
726                line_number: 0,
727                attributes: [("data".to_string(), "Hello < & World!".to_string())].into(),
728            }),
729            Item::Command(Command {
730                name: "tag".to_string(),
731                line_number: 0,
732                attributes: BTreeMap::new(),
733            }),
734            Item::Command(Command {
735                name: "tags".to_string(),
736                line_number: 0,
737                attributes: [("x".to_string(), "123".to_string())].into(),
738            }),
739            Item::Command(Command {
740                name: "name".to_string(),
741                line_number: 0,
742                attributes: [("0".to_string(), "Ok".to_string())].into(),
743            }),
744            Item::Command(Command {
745                name: "print".to_string(),
746                line_number: 0,
747                attributes: [("data".to_string(), "Test".to_string())].into(),
748            }),
749            Item::Command(Command {
750                name: "hcls".to_string(),
751                line_number: 0,
752                attributes: BTreeMap::from([("0".to_string(), "1".to_string())]),
753            }),
754        ]
755    )
756}